{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br/>\n",
    "\n",
    "# Part 3: Harness the Beast -- Fine-Tuning Google Bert & OpenAI GPT2  \n",
    "\n",
    "<br/>\n",
    "\n",
    "Transfer learning is one of the most important method to train a state-of-the-art NLP model after the Ulmfit Paper & Google Bert came out. Fine-tuning Bert & GPT2 requires huge computation power, but the conclusion is -- it totally worth it.\n",
    "\n",
    "There are several important things when fine-tuning Bert:\n",
    "\n",
    "## **1. translate pretrained weight<br/>**\n",
    "The Google Bert is trained by tensorflow framework. But we want to use pytorch to fine tune the model because it is more user friendly. Then we have to translate the tensorflow pretrained weight to pytorch pretrain weight. There's **pytorch_pretrained_bert** package out there help us to do so.\n",
    "\n",
    "\n",
    "## **2. Warm-up <br/>**\n",
    "The triangle shape learning rate schedule works well for fine-tuning (according the paper). It has a warm-up time (about 5% of total step, it's a hyper-parameter), the learning rate growth from zero to full learning rate at the point when warm-up is just over, then it decay slowly and uniformly and at the end of training step it just reach 0. So we need to calculate the number of step before we start training.\n",
    "\n",
    "## **3. Gradient Accumulations & Automatic Mixed Precision <br/>**\n",
    "Batch Size is important hyper-parameter in deep-learning. Batch size too small will lead to unstable gradient (and is hard to compute parallelly). But the problem is that in kaggle kernel the GPU Memory is only 16G (Tesla P100) not enough for even batch size 32. There's 2 approach to deal with this:\n",
    "\n",
    "**Automatic Mixed Precision**<br/>\n",
    "The great company **Nvidia** recently released a very useful package ---- **AMP**<br/>\n",
    "\n",
    "There are 3 main benefit:<br/>\n",
    "**Speeds up math-intensive operations**, such as linear and convolution layers, by using Tensor Cores.<br/>\n",
    "**Speeds up memory-limited operations** by accessing half the bytes compared to single-precision.<br/>\n",
    "**Reduces memory requirements for training models**, enabling larger models or larger minibatches.<br/>\n",
    "\n",
    "It's so useful and so easy to plug & play, so we have no reason to reject it! :)<br/>\n",
    "\n",
    "For detail explaination, please check [here](https://developer.nvidia.com/automatic-mixed-precision)\n",
    "\n",
    "\n",
    "**Gradient Accumulations**<br/>\n",
    "By using AMP the batch size can reach 32 in kaggle kernel, but still not enough.\n",
    "One more approach is gradient accumulation. At each step we just accumulate the gradient (without update the parameters), and when the batch_size*step is large enough (64 96 or 128) we just update the parameters, which mimic larger batch size.<br/>\n",
    "It's easy to understand but we need to be careful when implementing it, or the whole model may become garbage.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bert Fine-Tuning Pipline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "import numpy as np \n",
    "import pandas as pd\n",
    "from joblib import Parallel, delayed\n",
    "import os\n",
    "import datetime\n",
    "import pkg_resources\n",
    "import seaborn as sns\n",
    "import time\n",
    "import scipy.stats as stats\n",
    "import gc\n",
    "import re\n",
    "import operator \n",
    "import sys\n",
    "from sklearn import metrics\n",
    "from sklearn import model_selection\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data\n",
    "import torch.nn.functional as F\n",
    "from nltk.stem import PorterStemmer\n",
    "from sklearn.metrics import roc_auc_score\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "from tqdm import tqdm, tqdm_notebook\n",
    "import os\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = \"all\"\n",
    "import warnings\n",
    "warnings.filterwarnings(action='once')\n",
    "import pickle\n",
    "from apex import amp\n",
    "import shutil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.cuda.device_count() # AWS p3.8xlarge --- 4 Tesla V100 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1804874"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pd.read_csv('train.csv').__len__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQUENCE_LENGTH = 230\n",
    "SEED = 2025\n",
    "EPOCHS = 2\n",
    "Data_dir=\"./\"\n",
    "Input_dir = \"./\"  #\n",
    "WORK_DIR = \"./working/\" # Create by yourself to save pytorch Bert weight\n",
    "num_to_load=1704874   #Train Size\n",
    "valid_size= 100000   #Validation Size\n",
    "TOXICITY_COLUMN = 'target'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/importlib/_bootstrap.py:219: ImportWarning: can't resolve package from __spec__ or __package__, falling back on __name__ and __path__\n",
      "  return f(*args, **kwds)\n",
      "/home/ubuntu/.local/lib/python3.6/site-packages/tensorflow/python/keras/backend.py:5506: ResourceWarning: unclosed file <_io.TextIOWrapper name='/home/ubuntu/.keras/keras.json' mode='r' encoding='UTF-8'>\n",
      "  _config = json.load(open(_config_path))\n",
      "/home/ubuntu/anaconda3/envs/pytorch_p36/lib/python3.6/importlib/_bootstrap.py:219: RuntimeWarning: numpy.dtype size changed, may indicate binary incompatibility. Expected 96, got 88\n",
      "  return f(*args, **kwds)\n"
     ]
    }
   ],
   "source": [
    "from pytorch_pretrained_bert import convert_tf_checkpoint_to_pytorch\n",
    "from pytorch_pretrained_bert import BertTokenizer, BertForSequenceClassification,BertAdam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building PyTorch model from configuration: {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n",
      "Converting TensorFlow checkpoint from /home/ubuntu/.jupyter/uncased_L-12_H-768_A-12/bert_model.ckpt\n",
      "Loading TF weight bert/embeddings/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/embeddings/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/embeddings/position_embeddings with shape [512, 768]\n",
      "Loading TF weight bert/embeddings/token_type_embeddings with shape [2, 768]\n",
      "Loading TF weight bert/embeddings/word_embeddings with shape [30522, 768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_0/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_0/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_0/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_0/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_1/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_1/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_1/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_1/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_10/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_10/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_10/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_10/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_11/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_11/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_11/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_11/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_2/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_2/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_2/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_2/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_3/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_3/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_3/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_3/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_4/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_4/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_4/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_4/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_5/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_5/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_5/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_5/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_6/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_6/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_6/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_6/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_7/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_7/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_7/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_7/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_8/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_8/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_8/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_8/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/output/dense/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/key/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/key/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/query/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/query/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/value/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/attention/self/value/kernel with shape [768, 768]\n",
      "Loading TF weight bert/encoder/layer_9/intermediate/dense/bias with shape [3072]\n",
      "Loading TF weight bert/encoder/layer_9/intermediate/dense/kernel with shape [768, 3072]\n",
      "Loading TF weight bert/encoder/layer_9/output/LayerNorm/beta with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/output/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/output/dense/bias with shape [768]\n",
      "Loading TF weight bert/encoder/layer_9/output/dense/kernel with shape [3072, 768]\n",
      "Loading TF weight bert/pooler/dense/bias with shape [768]\n",
      "Loading TF weight bert/pooler/dense/kernel with shape [768, 768]\n",
      "Loading TF weight cls/predictions/output_bias with shape [30522]\n",
      "Loading TF weight cls/predictions/transform/LayerNorm/beta with shape [768]\n",
      "Loading TF weight cls/predictions/transform/LayerNorm/gamma with shape [768]\n",
      "Loading TF weight cls/predictions/transform/dense/bias with shape [768]\n",
      "Loading TF weight cls/predictions/transform/dense/kernel with shape [768, 768]\n",
      "Loading TF weight cls/seq_relationship/output_bias with shape [2]\n",
      "Loading TF weight cls/seq_relationship/output_weights with shape [2, 768]\n",
      "Initialize PyTorch weight ['bert', 'embeddings', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'embeddings', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'embeddings', 'position_embeddings']\n",
      "Initialize PyTorch weight ['bert', 'embeddings', 'token_type_embeddings']\n",
      "Initialize PyTorch weight ['bert', 'embeddings', 'word_embeddings']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_0', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_1', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_10', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_11', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_2', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_3', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_4', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_5', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_6', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_7', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_8', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'key', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'key', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'query', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'query', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'value', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'attention', 'self', 'value', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'intermediate', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'intermediate', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'output', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'output', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'output', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'encoder', 'layer_9', 'output', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['bert', 'pooler', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['bert', 'pooler', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['cls', 'predictions', 'output_bias']\n",
      "Initialize PyTorch weight ['cls', 'predictions', 'transform', 'LayerNorm', 'beta']\n",
      "Initialize PyTorch weight ['cls', 'predictions', 'transform', 'LayerNorm', 'gamma']\n",
      "Initialize PyTorch weight ['cls', 'predictions', 'transform', 'dense', 'bias']\n",
      "Initialize PyTorch weight ['cls', 'predictions', 'transform', 'dense', 'kernel']\n",
      "Initialize PyTorch weight ['cls', 'seq_relationship', 'output_bias']\n",
      "Initialize PyTorch weight ['cls', 'seq_relationship', 'output_weights']\n",
      "Save PyTorch model to ./working/pytorch_model.bin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'./working/bert_config.json'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Translate model from tensorflow to pytorch\n",
    "BERT_MODEL_PATH = './uncased_L-12_H-768_A-12/'\n",
    "convert_tf_checkpoint_to_pytorch.convert_tf_checkpoint_to_pytorch(\n",
    "    BERT_MODEL_PATH + 'bert_model.ckpt',\n",
    "BERT_MODEL_PATH + 'bert_config.json',\n",
    "WORK_DIR + 'pytorch_model.bin')\n",
    "\n",
    "shutil.copyfile(BERT_MODEL_PATH + 'bert_config.json', WORK_DIR + 'bert_config.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is the Bert configuration file\n",
    "from pytorch_pretrained_bert import BertConfig\n",
    "bert_config = BertConfig('./uncased_L-12_H-768_A-12/'+'bert_config.json')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Converting the lines to BERT format\n",
    "def convert_lines(example, max_seq_length,tokenizer):\n",
    "    max_seq_length -=2\n",
    "    all_tokens = []\n",
    "    longer = 0\n",
    "    for text in tqdm(example):\n",
    "        tokens_a = tokenizer.tokenize(text)\n",
    "        if len(tokens_a)>max_seq_length:\n",
    "            tokens_a = tokens_a[:max_seq_length]\n",
    "            longer += 1\n",
    "        one_token = tokenizer.convert_tokens_to_ids([\"[CLS]\"]+tokens_a+[\"[SEP]\"])+[0] * (max_seq_length - len(tokens_a))\n",
    "        all_tokens.append(one_token)\n",
    "    print(longer)\n",
    "    return np.array(all_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "identity_columns = [\n",
    "    'male', 'female', 'homosexual_gay_or_lesbian', 'christian', 'jewish',\n",
    "    'muslim', 'black', 'white', 'psychiatric_or_mental_illness']# for calculating validation score\n",
    "\n",
    "# for custom loss\n",
    "y_columns=['target']\n",
    "y_aux_columns = ['target_prob','target_prob','target_prob','severe_toxicity', 'obscene', 'identity_attack', 'insult', 'threat']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained(BERT_MODEL_PATH, cache_dir=None,do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loaded 1804874 records\n"
     ]
    }
   ],
   "source": [
    "train_df = pd.read_csv(os.path.join(Data_dir,\"train.csv\")).sample(num_to_load+valid_size,random_state=SEED)\n",
    "print('loaded %d records' % len(train_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1804874/1804874 [30:15<00:00, 994.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20250\n",
      "CPU times: user 30min 47s, sys: 4.94 s, total: 30min 52s\n",
      "Wall time: 30min 45s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Make sure all comment_text values are strings\n",
    "train_df['comment_text'] = train_df['comment_text'].astype(str) \n",
    "sequences = convert_lines(train_df[\"comment_text\"].fillna(\"DUMMY_VALUE\"),MAX_SEQUENCE_LENGTH,tokenizer)\n",
    "train_df=train_df.fillna(0)\n",
    "# List all identities\n",
    "\n",
    "\n",
    "train_df = train_df.drop(['comment_text'],axis=1)\n",
    "# convert target to 0,1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['target_prob']=train_df['target']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['target']=(train_df['target']>=0.5).astype(float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fw = open('sequences180_with_length','wb')  \n",
    "# pickle.dump(sequences, fw, -1)  \n",
    "# pickle.dump(seq_length,fw)\n",
    "# pickle.dump(train_df, fw)  \n",
    "# fw.close()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fw = open('sequences','rb')  \n",
    "# sequences = pickle.load(fw)\n",
    "# train_df = pickle.load(fw)  \n",
    "# fw.close()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# np.savez(\"180w_seed_2098_len_230.npz\", sequences)\n",
    "# train_df.to_csv('180w_seed_2098_len_230.csv',index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# r = np.load(\"180w_seed_2098_len_230.npz\")\n",
    "# sequences = r['arr_0']\n",
    "# train_df = pd.read_csv('180w_seed_2098_len_230.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1804874, 230)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sequences.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1804874, 45)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build up weight for custom loss\n",
    "\n",
    "# Overall\n",
    "weights = np.ones((len(train_df),)) / 4\n",
    "\n",
    "# Subgroup\n",
    "weights += (train_df[identity_columns].fillna(0).values>=0.5).sum(axis=1).astype(bool).astype(np.int) / 4\n",
    "\n",
    "# Background Positive, Subgroup Negative\n",
    "weights += (( (train_df['target'].values>=0.5).astype(bool).astype(np.int) +\n",
    "   (train_df[identity_columns].fillna(0).values<0.5).sum(axis=1).astype(bool).astype(np.int) ) > 1 ).astype(bool).astype(np.int) / 4\n",
    "\n",
    "# Background Negative, Subgroup Positive\n",
    "weights += (( (train_df['target'].values<0.5).astype(bool).astype(np.int) +\n",
    "   (train_df[identity_columns].fillna(0).values>=0.5).sum(axis=1).astype(bool).astype(np.int) ) > 1 ).astype(bool).astype(np.int) / 4\n",
    "loss_weight = 1.0 / weights.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = sequences[:num_to_load]                \n",
    "y = train_df[y_columns+y_aux_columns].values[:num_to_load]\n",
    "X_val = sequences[num_to_load:]                \n",
    "y_val = train_df[y_columns+y_aux_columns].values[num_to_load:]\n",
    "weights_train = weights[:num_to_load] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df=train_df.tail(valid_size).copy()\n",
    "train_df=train_df.head(num_to_load)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f69d70d65f0>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output_model_file = \"2epoch_no_dense_170w_seqbuck.bin\"\n",
    "lr=2e-5\n",
    "batch_size = 128\n",
    "accumulation_steps=1\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_torch = torch.tensor(weights_train,dtype=torch.float)\n",
    "train_dataset = torch.utils.data.TensorDataset(torch.tensor(X,dtype=torch.long),torch.tensor(y,dtype=torch.float),weights_torch)\n",
    "train = train_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Buildup Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "BERT_OUT = 128\n",
    "DENSE_HIDDEN_UNITS = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNet, self).__init__()\n",
    "        self.bert_layer = BertForSequenceClassification.from_pretrained(\"./working\",cache_dir=None,num_labels = BERT_OUT)\n",
    "        self.linear_out = nn.Linear(BERT_OUT, 1)\n",
    "        self.linear_aux_out = nn.Linear(BERT_OUT, 8)\n",
    "        #self.identity_out = nn.Linear(BERT_OUT, 9)\n",
    "        self.drop_out_layer = nn.Dropout(p=0.1)        \n",
    "        \n",
    "    def forward(self, x, attention_mask=None, labels=None):\n",
    "        bert_out = self.bert_layer(x, attention_mask=attention_mask, labels=labels)\n",
    "        drop_out_layer_out = self.drop_out_layer(bert_out)\n",
    "        result = self.linear_out(drop_out_layer_out)\n",
    "        aux_result = self.linear_aux_out(bert_out)\n",
    "        #identity_result = self.identity_out(bert_out)\n",
    "        out = torch.cat([result, aux_result], 1)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NeuralNet()\n",
    "model.zero_grad()\n",
    "model = model.cuda()\n",
    "param_optimizer = list(model.named_parameters())\n",
    "no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "26638"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# calculate the step that learning rate reach 0\n",
    "num_train_optimization_steps = int(EPOCHS*len(train)/batch_size/accumulation_steps)\n",
    "num_train_optimization_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# original Bert setting\n",
    "optimizer_grouped_parameters = [\n",
    "    {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},\n",
    "    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}\n",
    "    ]\n",
    "\n",
    "# warmup is important\n",
    "optimizer = BertAdam(optimizer_grouped_parameters,\n",
    "                     lr=lr,\n",
    "                     warmup=0.05,\n",
    "                     t_total=num_train_optimization_steps)\n",
    "\n",
    "model, optimizer = amp.initialize(model, optimizer, opt_level=\"O1\",verbosity=0)\n",
    "\n",
    "model=model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.nn.DataParallel(model) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Number of Trainerable Parameters Group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "79"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer_grouped_parameters[0]['params'].__len__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "126"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer_grouped_parameters[1]['params'].__len__()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loss Fuction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_loss(preds,targets,weights):\n",
    "    ''' Define custom loss function for weighted BCE on 'target' column '''\n",
    "    bce_loss_1 = nn.BCEWithLogitsLoss(weight=weights)(preds[:,0],targets[:,0])\n",
    "    bce_loss_2 = nn.BCEWithLogitsLoss()(preds[:,1:],targets[:,1:])\n",
    "    return ((bce_loss_1 * loss_weight)*0.60 + bce_loss_2*0.40)*2 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 100\n",
      "time_elapsed: 1.4143638451894125 min 1.1232507607062325\n",
      "step: 200\n",
      "time_elapsed: 2.6277382373809814 min 0.8408596892799032\n",
      "step: 300\n",
      "time_elapsed: 3.8333160916964215 min 0.6025500591032656\n",
      "step: 400\n",
      "time_elapsed: 5.036450576782227 min 0.4747738666715351\n",
      "step: 500\n",
      "time_elapsed: 6.239763224124909 min 0.4009440955322863\n",
      "step: 600\n",
      "time_elapsed: 7.442309661706289 min 0.3805666459231075\n",
      "step: 700\n",
      "time_elapsed: 8.650884234905243 min 0.3692055115898399\n",
      "step: 800\n",
      "time_elapsed: 9.861127376556396 min 0.35607749229518393\n",
      "step: 900\n",
      "time_elapsed: 11.067019267876942 min 0.3506028733813989\n",
      "step: 1000\n",
      "time_elapsed: 12.269352165857951 min 0.35626946913882585\n",
      "step: 1100\n",
      "time_elapsed: 13.472266602516175 min 0.3469740949427739\n",
      "step: 1200\n",
      "time_elapsed: 14.675750207901 min 0.3444224798299555\n",
      "step: 1300\n",
      "time_elapsed: 15.879150104522704 min 0.3323961508100495\n",
      "step: 1400\n",
      "time_elapsed: 17.088626070817313 min 0.33060169265859163\n",
      "step: 1500\n",
      "time_elapsed: 18.293953319390614 min 0.3346307585501615\n",
      "step: 1600\n",
      "time_elapsed: 19.496057720979056 min 0.32536484390228215\n",
      "step: 1700\n",
      "time_elapsed: 20.69234098990758 min 0.3460572769286651\n",
      "step: 1800\n",
      "time_elapsed: 21.894945216178893 min 0.33231555552661235\n",
      "step: 1900\n",
      "time_elapsed: 23.095089852809906 min 0.3348228102752444\n",
      "step: 2000\n",
      "time_elapsed: 24.288453908761344 min 0.32867358815595954\n",
      "step: 2100\n",
      "time_elapsed: 25.475119654337565 min 0.32504156503983284\n",
      "step: 2200\n",
      "time_elapsed: 26.661754250526428 min 0.3287601012634513\n",
      "step: 2300\n",
      "time_elapsed: 27.85670292377472 min 0.3354028266091158\n",
      "step: 2400\n",
      "time_elapsed: 29.048945526281994 min 0.3123845924517165\n",
      "step: 2500\n",
      "time_elapsed: 30.237513677279154 min 0.30678890492741673\n",
      "step: 2600\n",
      "time_elapsed: 31.43377762635549 min 0.32076828979320754\n",
      "step: 2700\n",
      "time_elapsed: 32.62868094444275 min 0.32143379903996605\n",
      "step: 2800\n",
      "time_elapsed: 33.821103934446974 min 0.327494494408679\n",
      "step: 2900\n",
      "time_elapsed: 35.01590898831685 min 0.33036168374418384\n",
      "step: 3000\n",
      "time_elapsed: 36.21211139758428 min 0.31761177542595886\n",
      "step: 3100\n",
      "time_elapsed: 37.40417688290278 min 0.31451411820744635\n",
      "step: 3200\n",
      "time_elapsed: 38.60475399096807 min 0.3154262667676901\n",
      "step: 3300\n",
      "time_elapsed: 39.80127039750417 min 0.31548017538344725\n",
      "step: 3400\n",
      "time_elapsed: 41.00659199555715 min 0.3205221045068641\n",
      "step: 3500\n",
      "time_elapsed: 42.20819576978683 min 0.3093250936931484\n",
      "step: 3600\n",
      "time_elapsed: 43.41006304423014 min 0.3193544508548859\n",
      "step: 3700\n",
      "time_elapsed: 44.60390493075053 min 0.30963236210995226\n",
      "step: 3800\n",
      "time_elapsed: 45.7951450030009 min 0.3080822760576591\n",
      "step: 3900\n",
      "time_elapsed: 46.99099044005076 min 0.3144059779057393\n",
      "step: 4000\n",
      "time_elapsed: 48.19090827703476 min 0.31111952532163634\n",
      "step: 4100\n",
      "time_elapsed: 49.380025959014894 min 0.3261672053709443\n",
      "step: 4200\n",
      "time_elapsed: 50.5739432533582 min 0.3215537785464688\n",
      "step: 4300\n",
      "time_elapsed: 51.776288612683615 min 0.3279641720559248\n",
      "step: 4400\n",
      "time_elapsed: 52.97558421293895 min 0.3180222405230487\n",
      "step: 4500\n",
      "time_elapsed: 54.17425269285838 min 0.3039881907610884\n",
      "step: 4600\n",
      "time_elapsed: 55.36739500761032 min 0.30978463136164364\n",
      "step: 4700\n",
      "time_elapsed: 56.54689562320709 min 0.3174176551179772\n",
      "step: 4800\n",
      "time_elapsed: 57.72512957255046 min 0.3270832281154534\n",
      "step: 4900\n",
      "time_elapsed: 58.90717920064926 min 0.30327266289046434\n",
      "step: 5000\n",
      "time_elapsed: 60.08717488447825 min 0.3141813926306135\n",
      "step: 5100\n",
      "time_elapsed: 61.281583746274315 min 0.3145815003302341\n",
      "step: 5200\n",
      "time_elapsed: 62.47661258776983 min 0.32381412844720914\n",
      "step: 5300\n",
      "time_elapsed: 63.67426374355952 min 0.3118638740614951\n",
      "step: 5400\n",
      "time_elapsed: 64.86640360752742 min 0.3018743441037563\n",
      "step: 5500\n",
      "time_elapsed: 66.06460750500361 min 0.3112861167194945\n",
      "step: 5600\n",
      "time_elapsed: 67.25075764656067 min 0.30693970245306784\n",
      "step: 5700\n",
      "time_elapsed: 68.440500831604 min 0.3098133192190479\n",
      "step: 5800\n",
      "time_elapsed: 69.62975774208705 min 0.318459669306288\n",
      "step: 5900\n",
      "time_elapsed: 70.81898531516393 min 0.30984470218592136\n",
      "step: 6000\n",
      "time_elapsed: 72.00621773004532 min 0.3056060479139874\n",
      "step: 6100\n",
      "time_elapsed: 73.20279420216879 min 0.3074840223774071\n",
      "step: 6200\n",
      "time_elapsed: 74.39492862224579 min 0.3091336564920831\n",
      "step: 6300\n",
      "time_elapsed: 75.59204729795457 min 0.3088937817514319\n",
      "step: 6400\n",
      "time_elapsed: 76.79046114285786 min 0.3167659766341713\n",
      "step: 6500\n",
      "time_elapsed: 77.99067112604777 min 0.31149363245439576\n",
      "step: 6600\n",
      "time_elapsed: 79.1919368426005 min 0.30770966144522705\n",
      "step: 6700\n",
      "time_elapsed: 80.39143945376078 min 0.30788532962415677\n",
      "step: 6800\n",
      "time_elapsed: 81.5901108344396 min 0.30463546425279503\n",
      "step: 6900\n",
      "time_elapsed: 82.7975121418635 min 0.300967596063363\n",
      "step: 7000\n",
      "time_elapsed: 84.00124475161235 min 0.30397018327429143\n",
      "step: 7100\n",
      "time_elapsed: 85.19986984332402 min 0.29992331726367116\n",
      "step: 7200\n",
      "time_elapsed: 86.40148155689239 min 0.2900259361467367\n",
      "step: 7300\n",
      "time_elapsed: 87.60573307275772 min 0.2872336756496114\n",
      "step: 7400\n",
      "time_elapsed: 88.81276840368906 min 0.3157991847020631\n",
      "step: 7500\n",
      "time_elapsed: 90.01609164476395 min 0.31214712705586295\n",
      "step: 7600\n",
      "time_elapsed: 91.21911184390386 min 0.31847273649828567\n",
      "step: 7700\n",
      "time_elapsed: 92.42391907771429 min 0.31527571170465973\n",
      "step: 7800\n",
      "time_elapsed: 93.62666341463725 min 0.3049247038001075\n",
      "step: 7900\n",
      "time_elapsed: 94.827132999897 min 0.3046727287035377\n",
      "step: 8000\n",
      "time_elapsed: 96.02581355969112 min 0.31178015104177986\n",
      "step: 8100\n",
      "time_elapsed: 97.23382654984792 min 0.30996041137278996\n",
      "step: 8200\n",
      "time_elapsed: 98.4387156844139 min 0.3018689825113205\n",
      "step: 8300\n",
      "time_elapsed: 99.63655378818513 min 0.30562708571603325\n",
      "step: 8400\n",
      "time_elapsed: 100.84258980751038 min 0.3033985653596515\n",
      "step: 8500\n",
      "time_elapsed: 102.04587310949961 min 0.30772697001108096\n",
      "step: 8600\n",
      "time_elapsed: 103.25154681603114 min 0.30441337819517783\n",
      "step: 8700\n",
      "time_elapsed: 104.45384342273077 min 0.31331663789337355\n",
      "step: 8800\n",
      "time_elapsed: 105.6523153146108 min 0.3019879450476862\n",
      "step: 8900\n",
      "time_elapsed: 106.8569440484047 min 0.3043738528039004\n",
      "step: 9000\n",
      "time_elapsed: 108.06568905909856 min 0.29878460149172115\n",
      "step: 9100\n",
      "time_elapsed: 109.2656267841657 min 0.3044192202926718\n",
      "step: 9200\n",
      "time_elapsed: 110.47616861263911 min 0.2951538230986289\n",
      "step: 9300\n",
      "time_elapsed: 111.67404284874598 min 0.31733888137778615\n",
      "step: 9400\n",
      "time_elapsed: 112.87836982409159 min 0.3012885450335073\n",
      "step: 9500\n",
      "time_elapsed: 114.07480691274007 min 0.291704689297812\n",
      "step: 9600\n",
      "time_elapsed: 115.27874969641367 min 0.28730156454505745\n",
      "step: 9700\n",
      "time_elapsed: 116.4752825975418 min 0.301387612024126\n",
      "step: 9800\n",
      "time_elapsed: 117.67876015504201 min 0.29933958262257765\n",
      "step: 9900\n",
      "time_elapsed: 118.88481553395589 min 0.30789019036771587\n",
      "step: 10000\n",
      "time_elapsed: 120.09486794471741 min 0.30372732943285224\n",
      "step: 10100\n",
      "time_elapsed: 121.29921496311823 min 0.30051319715409175\n",
      "step: 10200\n",
      "time_elapsed: 122.5055813550949 min 0.3065406518996307\n",
      "step: 10300\n",
      "time_elapsed: 123.71372106870015 min 0.3019529669090383\n",
      "step: 10400\n",
      "time_elapsed: 124.92054181496302 min 0.29643568914797386\n",
      "step: 10500\n",
      "time_elapsed: 126.13166220585505 min 0.30182033732720276\n",
      "step: 10600\n",
      "time_elapsed: 127.33806236187617 min 0.29907369859417854\n",
      "step: 10700\n",
      "time_elapsed: 128.54484455982845 min 0.2945596112746211\n",
      "step: 10800\n",
      "time_elapsed: 129.75054406325023 min 0.27878818368966346\n",
      "step: 10900\n",
      "time_elapsed: 130.9542415102323 min 0.30921916924006804\n",
      "step: 11000\n",
      "time_elapsed: 132.1605100909869 min 0.30805998173280713\n",
      "step: 11100\n",
      "time_elapsed: 133.36670002142588 min 0.30200192110532287\n",
      "step: 11200\n",
      "time_elapsed: 134.56888707081478 min 0.30396679936070503\n",
      "step: 11300\n",
      "time_elapsed: 135.77399453719457 min 0.306216571832405\n",
      "step: 11400\n",
      "time_elapsed: 136.98049103418987 min 0.30591991451729306\n",
      "step: 11500\n",
      "time_elapsed: 138.18338627020518 min 0.2935704886122685\n",
      "step: 11600\n",
      "time_elapsed: 139.39391920963922 min 0.2982649020765628\n",
      "step: 11700\n",
      "time_elapsed: 140.60090488592783 min 0.2943971174399064\n",
      "step: 11800\n",
      "time_elapsed: 141.81260633468628 min 0.3064950419461693\n",
      "step: 11900\n",
      "time_elapsed: 143.02049913406373 min 0.301409700066363\n",
      "step: 12000\n",
      "time_elapsed: 144.22677264610925 min 0.29659131318648196\n",
      "step: 12100\n",
      "time_elapsed: 145.42761714458464 min 0.30282220985807706\n",
      "step: 12200\n",
      "time_elapsed: 146.6272980093956 min 0.30734919867760124\n",
      "step: 12300\n",
      "time_elapsed: 147.83121682008107 min 0.2917704877163546\n",
      "step: 12400\n",
      "time_elapsed: 149.03212426900865 min 0.2918082832283186\n",
      "step: 12500\n",
      "time_elapsed: 150.23588880697886 min 0.2851359742549588\n",
      "step: 12600\n",
      "time_elapsed: 151.43906679153443 min 0.30143756965493507\n",
      "step: 12700\n",
      "time_elapsed: 152.64440871874493 min 0.29530227112772556\n",
      "step: 12800\n",
      "time_elapsed: 153.84533838033676 min 0.3059240389034935\n",
      "step: 12900\n",
      "time_elapsed: 155.0496228257815 min 0.309400802815155\n",
      "step: 13000\n",
      "time_elapsed: 156.25424469709395 min 0.29999900342585306\n",
      "step: 13100\n",
      "time_elapsed: 157.45653829574584 min 0.2883600045282188\n",
      "step: 13200\n",
      "time_elapsed: 158.65928637981415 min 0.2953965786884564\n",
      "step: 13300\n",
      "time_elapsed: 159.85597836971283 min 0.30685119558052926\n",
      "step: 100\n",
      "time_elapsed: 1.2043689807256064 min 0.27307855596798997\n",
      "step: 200\n",
      "time_elapsed: 2.4020887931187946 min 0.26744760019307345\n",
      "step: 300\n",
      "time_elapsed: 3.598984718322754 min 0.28167117032641026\n",
      "step: 400\n",
      "time_elapsed: 4.8047760963439945 min 0.26689015516549147\n",
      "step: 500\n",
      "time_elapsed: 6.007501105467479 min 0.2780328832163286\n",
      "step: 600\n",
      "time_elapsed: 7.21170080502828 min 0.274415245531829\n",
      "step: 700\n",
      "time_elapsed: 8.415684223175049 min 0.2672685576640017\n",
      "step: 800\n",
      "time_elapsed: 9.616731107234955 min 0.2734148315312077\n",
      "step: 900\n",
      "time_elapsed: 10.818064709504446 min 0.278360825043349\n",
      "step: 1000\n",
      "time_elapsed: 12.019630138079325 min 0.2761467946880856\n",
      "step: 1100\n",
      "time_elapsed: 13.224563896656036 min 0.2803839070909907\n",
      "step: 1200\n",
      "time_elapsed: 14.43152525027593 min 0.2686017068077712\n",
      "step: 1300\n",
      "time_elapsed: 15.633093551794689 min 0.26586192540585313\n",
      "step: 1400\n",
      "time_elapsed: 16.84041386047999 min 0.2721520775658873\n",
      "step: 1500\n",
      "time_elapsed: 18.046548322836557 min 0.25665666955698313\n",
      "step: 1600\n",
      "time_elapsed: 19.254971277713775 min 0.26582773169017443\n",
      "step: 1700\n",
      "time_elapsed: 20.46298764149348 min 0.2833942038708457\n",
      "step: 1800\n",
      "time_elapsed: 21.669161784648896 min 0.27698708696225177\n",
      "step: 1900\n",
      "time_elapsed: 22.86889796257019 min 0.2801708665051827\n",
      "step: 2000\n",
      "time_elapsed: 24.07052748600642 min 0.2726737258692583\n",
      "step: 2100\n",
      "time_elapsed: 25.271219849586487 min 0.2694940849425234\n",
      "step: 2200\n",
      "time_elapsed: 26.466856356461843 min 0.2632989278853568\n",
      "step: 2300\n",
      "time_elapsed: 27.664367620150248 min 0.2739146906898452\n",
      "step: 2400\n",
      "time_elapsed: 28.859725741545358 min 0.27872436499062536\n",
      "step: 2500\n",
      "time_elapsed: 30.045994587739308 min 0.26977823612500484\n",
      "step: 2600\n",
      "time_elapsed: 31.241398294766743 min 0.27017321607819716\n",
      "step: 2700\n",
      "time_elapsed: 32.42924913565318 min 0.2709034964624891\n",
      "step: 2800\n",
      "time_elapsed: 33.62130638758342 min 0.27705846261867073\n",
      "step: 2900\n",
      "time_elapsed: 34.81698884963989 min 0.27144636153332036\n",
      "step: 3000\n",
      "time_elapsed: 36.00139742294947 min 0.279839507942813\n",
      "step: 3100\n",
      "time_elapsed: 37.19279188712438 min 0.27667903709243774\n",
      "step: 3200\n",
      "time_elapsed: 38.39407987594605 min 0.27429089560474346\n",
      "step: 3300\n",
      "time_elapsed: 39.585775049527484 min 0.2741617273206018\n",
      "step: 3400\n",
      "time_elapsed: 40.781333939234415 min 0.2741075012939963\n",
      "step: 3500\n",
      "time_elapsed: 41.97913136084875 min 0.2713895910113408\n",
      "step: 3600\n",
      "time_elapsed: 43.18006615241369 min 0.2825885348646276\n",
      "step: 3700\n",
      "time_elapsed: 44.386003037293754 min 0.2684900145931024\n",
      "step: 3800\n",
      "time_elapsed: 45.58944929043452 min 0.27056583124460903\n",
      "step: 3900\n",
      "time_elapsed: 46.78283730745316 min 0.27336033750695826\n",
      "step: 4000\n",
      "time_elapsed: 47.97964193820953 min 0.27074589389840376\n",
      "step: 4100\n",
      "time_elapsed: 49.17590324878692 min 0.273140162157877\n",
      "step: 4200\n",
      "time_elapsed: 50.36459271510442 min 0.27069133058247274\n",
      "step: 4300\n",
      "time_elapsed: 51.55818114678065 min 0.27660676643972204\n",
      "step: 4400\n",
      "time_elapsed: 52.7560244957606 min 0.26596731150614417\n",
      "step: 4500\n",
      "time_elapsed: 53.955435852209725 min 0.2808204566831871\n",
      "step: 4600\n",
      "time_elapsed: 55.159919309616086 min 0.2764007815262518\n",
      "step: 4700\n",
      "time_elapsed: 56.355564431349435 min 0.2729350772568133\n",
      "step: 4800\n",
      "time_elapsed: 57.55237024227778 min 0.26873030437513395\n",
      "step: 4900\n",
      "time_elapsed: 58.75060486793518 min 0.2756053823722753\n",
      "step: 5000\n",
      "time_elapsed: 59.949988508224486 min 0.266723205469372\n",
      "step: 5100\n",
      "time_elapsed: 61.14902586539586 min 0.27480261685053275\n",
      "step: 5200\n",
      "time_elapsed: 62.34495362440745 min 0.2780170568170295\n",
      "step: 5300\n",
      "time_elapsed: 63.542519819736484 min 0.2625808532990213\n",
      "step: 5400\n",
      "time_elapsed: 64.73875276645025 min 0.2714864306275404\n",
      "step: 5500\n",
      "time_elapsed: 65.93143160740534 min 0.2729175631492759\n",
      "step: 5600\n",
      "time_elapsed: 67.12461225589117 min 0.2744047712122521\n",
      "step: 5700\n",
      "time_elapsed: 68.31790894667307 min 0.25912214365694086\n",
      "step: 5800\n",
      "time_elapsed: 69.50770511627198 min 0.27286348787205134\n",
      "step: 5900\n",
      "time_elapsed: 70.70809683799743 min 0.2731234077233713\n",
      "step: 6000\n",
      "time_elapsed: 71.90386197964351 min 0.2755254917102825\n",
      "step: 6100\n",
      "time_elapsed: 73.09709570805232 min 0.27080946951438295\n",
      "step: 6200\n",
      "time_elapsed: 74.2908154686292 min 0.26807378293315953\n",
      "step: 6300\n",
      "time_elapsed: 75.47994658152263 min 0.27052226610047964\n",
      "step: 6400\n",
      "time_elapsed: 76.66611124674479 min 0.2631274393555819\n",
      "step: 6500\n",
      "time_elapsed: 77.85850387414297 min 0.27430149632455864\n",
      "step: 6600\n",
      "time_elapsed: 79.05629872878393 min 0.2724735133362314\n",
      "step: 6700\n",
      "time_elapsed: 80.2564690152804 min 0.2604077815288474\n",
      "step: 6800\n",
      "time_elapsed: 81.45124596357346 min 0.27392498864129466\n",
      "step: 6900\n",
      "time_elapsed: 82.6420942902565 min 0.2580503274057048\n",
      "step: 7000\n",
      "time_elapsed: 83.83798518975576 min 0.2550989192410663\n",
      "step: 7100\n",
      "time_elapsed: 85.03087642987569 min 0.2652519967151139\n",
      "step: 7200\n",
      "time_elapsed: 86.2235247651736 min 0.2638033676306034\n",
      "step: 7300\n",
      "time_elapsed: 87.41030659675599 min 0.2716964932845014\n",
      "step: 7400\n",
      "time_elapsed: 88.59804547230402 min 0.26712857814046886\n",
      "step: 7500\n",
      "time_elapsed: 89.78311143318813 min 0.27152784332720753\n",
      "step: 7600\n",
      "time_elapsed: 90.96979025999705 min 0.2643991482348503\n",
      "step: 7700\n",
      "time_elapsed: 92.15695619185766 min 0.2653839112393694\n",
      "step: 7800\n",
      "time_elapsed: 93.34089233875275 min 0.2745811079335343\n",
      "step: 7900\n",
      "time_elapsed: 94.5292701045672 min 0.2700310282128662\n",
      "step: 8000\n",
      "time_elapsed: 95.72008223930995 min 0.2589806408103256\n",
      "step: 8100\n",
      "time_elapsed: 96.90679117838542 min 0.2692391787568228\n",
      "step: 8200\n",
      "time_elapsed: 98.09694947004319 min 0.2617521639988987\n",
      "step: 8300\n",
      "time_elapsed: 99.2847882548968 min 0.2793411353922762\n",
      "step: 8400\n",
      "time_elapsed: 100.46474143266678 min 0.2727484856148526\n",
      "step: 8500\n",
      "time_elapsed: 101.66045885880789 min 0.2762963565826509\n",
      "step: 8600\n",
      "time_elapsed: 102.8529906908671 min 0.2643183154401448\n",
      "step: 8700\n",
      "time_elapsed: 104.0515918135643 min 0.2715110406725431\n",
      "step: 8800\n",
      "time_elapsed: 105.25241769154867 min 0.2654992567406967\n",
      "step: 8900\n",
      "time_elapsed: 106.45546662807465 min 0.27240046146896135\n",
      "step: 9000\n",
      "time_elapsed: 107.64859847625097 min 0.26167074594238443\n",
      "step: 9100\n",
      "time_elapsed: 108.84777584870656 min 0.2602993014270204\n",
      "step: 9200\n",
      "time_elapsed: 110.04465043544769 min 0.27835971008137944\n",
      "step: 9300\n",
      "time_elapsed: 111.23969389994939 min 0.2757641407733184\n",
      "step: 9400\n",
      "time_elapsed: 112.42985557715097 min 0.26831961285524036\n",
      "step: 9500\n",
      "time_elapsed: 113.61550621589025 min 0.26801389784357843\n",
      "step: 9600\n",
      "time_elapsed: 114.81005461613337 min 0.26819493064053535\n",
      "step: 9700\n",
      "time_elapsed: 115.99274537563323 min 0.2748168510005117\n",
      "step: 9800\n",
      "time_elapsed: 117.18719380696615 min 0.2641119895012914\n",
      "step: 9900\n",
      "time_elapsed: 118.37627373139064 min 0.25232176834751924\n",
      "step: 10000\n",
      "time_elapsed: 119.56618960301081 min 0.2564450696851358\n",
      "step: 10100\n",
      "time_elapsed: 120.7616523941358 min 0.26043673164873626\n",
      "step: 10200\n",
      "time_elapsed: 121.95979650815327 min 0.2626449331970896\n",
      "step: 10300\n",
      "time_elapsed: 123.15666601657867 min 0.2745077179169525\n",
      "step: 10400\n",
      "time_elapsed: 124.35371789534886 min 0.269714642262743\n",
      "step: 10500\n",
      "time_elapsed: 125.54992446104686 min 0.2638852664090766\n",
      "step: 10600\n",
      "time_elapsed: 126.7500896970431 min 0.2510430234563357\n",
      "step: 10700\n",
      "time_elapsed: 127.95143176317215 min 0.2590593238023714\n",
      "step: 10800\n",
      "time_elapsed: 129.13771723508836 min 0.2647002750884125\n",
      "step: 10900\n",
      "time_elapsed: 130.3325725078583 min 0.25862254030077053\n",
      "step: 11000\n",
      "time_elapsed: 131.51866259972255 min 0.2839149714191532\n",
      "step: 11100\n",
      "time_elapsed: 132.70650657812755 min 0.27642689149928856\n",
      "step: 11200\n",
      "time_elapsed: 133.89442214171092 min 0.26408053690267524\n",
      "step: 11300\n",
      "time_elapsed: 135.0867569764455 min 0.2650831505022568\n",
      "step: 11400\n",
      "time_elapsed: 136.2803432861964 min 0.2724148326983127\n",
      "step: 11500\n",
      "time_elapsed: 137.47609295050304 min 0.2665765193740848\n",
      "step: 11600\n",
      "time_elapsed: 138.6709841688474 min 0.2663616072730038\n",
      "step: 11700\n",
      "time_elapsed: 139.87244447867076 min 0.2558133211938758\n",
      "step: 11800\n",
      "time_elapsed: 141.06481379667918 min 0.26122210650214667\n",
      "step: 11900\n",
      "time_elapsed: 142.25956190824508 min 0.26513498670919444\n",
      "step: 12000\n",
      "time_elapsed: 143.46232268015544 min 0.2643295998325957\n",
      "step: 12100\n",
      "time_elapsed: 144.66342820326489 min 0.26876845710399627\n",
      "step: 12200\n",
      "time_elapsed: 145.8565006971359 min 0.27599411152220127\n",
      "step: 12300\n",
      "time_elapsed: 147.06301932732265 min 0.2668517166439483\n",
      "step: 12400\n",
      "time_elapsed: 148.26923031806945 min 0.264122092857234\n",
      "step: 12500\n",
      "time_elapsed: 149.4700691739718 min 0.2698484681603228\n",
      "step: 12600\n",
      "time_elapsed: 150.65730812152228 min 0.25981953591086465\n",
      "step: 12700\n",
      "time_elapsed: 151.8432366291682 min 0.2617634074912942\n",
      "step: 12800\n",
      "time_elapsed: 153.02980913321178 min 0.25956001540995227\n",
      "step: 12900\n",
      "time_elapsed: 154.21696432828904 min 0.264721706884314\n",
      "step: 13000\n",
      "time_elapsed: 155.40804793834687 min 0.2555947888774479\n",
      "step: 13100\n",
      "time_elapsed: 156.6056671222051 min 0.2636418852318957\n",
      "step: 13200\n",
      "time_elapsed: 157.79650489091873 min 0.2650578409953612\n",
      "step: 13300\n",
      "time_elapsed: 158.98299668629963 min 0.2644481496788835\n"
     ]
    }
   ],
   "source": [
    "tq = range(EPOCHS)\n",
    "for epoch in tq:\n",
    "    start_time = time.time()\n",
    "    train_loader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True)\n",
    "    #train_loader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True)\n",
    "    avg_loss = 0.\n",
    "    avg_accuracy = 0.\n",
    "    lossf=None\n",
    "    count = 0\n",
    "    optimizer.zero_grad()   # Bug fix - thanks to @chinhuic\n",
    "    for i,(x_batch, y_batch,weight_batch) in enumerate(train_loader):\n",
    "#        optimizer.zero_grad()\n",
    "        y_pred = model(x_batch.cuda(), attention_mask=(x_batch>0).cuda(), labels=None)\n",
    "        loss =  custom_loss(y_pred,y_batch.cuda(),weight_batch.cuda())\n",
    "        with amp.scale_loss(loss, optimizer) as scaled_loss:\n",
    "            scaled_loss.backward()\n",
    "        if (i+1) % accumulation_steps == 0:     # Wait for several backward steps\n",
    "            optimizer.step()                    # Now we can do an optimizer step\n",
    "            optimizer.zero_grad()    \n",
    "        if lossf:\n",
    "            lossf = 0.98*lossf+0.02*loss.item()\n",
    "        else:\n",
    "            lossf = loss.item()\n",
    "        count += 1\n",
    "        if count%100 == 0:\n",
    "            print('step:',count)\n",
    "            print('time_elapsed:',(time.time()-start_time)/60,'min',lossf)\n",
    "    torch.save(model.state_dict(), 'batch128_probtarget2_'+str(epoch)+'.bin')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "IncompatibleKeys(missing_keys=[], unexpected_keys=[])"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): NeuralNet(\n",
       "    (bert_layer): BertForSequenceClassification(\n",
       "      (bert): BertModel(\n",
       "        (embeddings): BertEmbeddings(\n",
       "          (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "          (position_embeddings): Embedding(512, 768)\n",
       "          (token_type_embeddings): Embedding(2, 768)\n",
       "          (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1)\n",
       "        )\n",
       "        (encoder): BertEncoder(\n",
       "          (layer): ModuleList(\n",
       "            (0): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (1): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (2): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (3): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (4): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (5): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (6): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (7): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (8): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (9): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (10): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (11): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (pooler): BertPooler(\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (activation): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (dropout): Dropout(p=0.1)\n",
       "      (classifier): Linear(in_features=768, out_features=128, bias=True)\n",
       "    )\n",
       "    (linear_out): Linear(in_features=128, out_features=1, bias=True)\n",
       "    (linear_aux_out): Linear(in_features=128, out_features=8, bias=True)\n",
       "    (drop_out_layer): Dropout(p=0.1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "DataParallel(\n",
       "  (module): NeuralNet(\n",
       "    (bert_layer): BertForSequenceClassification(\n",
       "      (bert): BertModel(\n",
       "        (embeddings): BertEmbeddings(\n",
       "          (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
       "          (position_embeddings): Embedding(512, 768)\n",
       "          (token_type_embeddings): Embedding(2, 768)\n",
       "          (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1)\n",
       "        )\n",
       "        (encoder): BertEncoder(\n",
       "          (layer): ModuleList(\n",
       "            (0): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (1): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (2): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (3): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (4): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (5): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (6): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (7): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (8): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (9): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (10): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "            (11): BertLayer(\n",
       "              (attention): BertAttention(\n",
       "                (self): BertSelfAttention(\n",
       "                  (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "                (output): BertSelfOutput(\n",
       "                  (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "                  (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                  (dropout): Dropout(p=0.1)\n",
       "                )\n",
       "              )\n",
       "              (intermediate): BertIntermediate(\n",
       "                (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "              )\n",
       "              (output): BertOutput(\n",
       "                (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "                (LayerNorm): FusedLayerNorm(torch.Size([768]), eps=1e-12, elementwise_affine=True)\n",
       "                (dropout): Dropout(p=0.1)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "        (pooler): BertPooler(\n",
       "          (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "          (activation): Tanh()\n",
       "        )\n",
       "      )\n",
       "      (dropout): Dropout(p=0.1)\n",
       "      (classifier): Linear(in_features=768, out_features=128, bias=True)\n",
       "    )\n",
       "    (linear_out): Linear(in_features=128, out_features=1, bias=True)\n",
       "    (linear_aux_out): Linear(in_features=128, out_features=8, bias=True)\n",
       "    (drop_out_layer): Dropout(p=0.1)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [03:28<00:00,  4.22it/s]\n"
     ]
    }
   ],
   "source": [
    "# Run validation\n",
    "model =NeuralNet()\n",
    "model = torch.nn.DataParallel(model)\n",
    "model.load_state_dict(torch.load('batch128_probtarget2_1.bin'))\n",
    "model.cuda()\n",
    "for param in model.parameters():\n",
    "    param.requires_grad=False\n",
    "model.eval()\n",
    "valid_preds = np.zeros((len(X_val)))\n",
    "valid = torch.utils.data.TensorDataset(torch.tensor(X_val,dtype=torch.long))\n",
    "valid_loader = torch.utils.data.DataLoader(valid, batch_size=128, shuffle=False)\n",
    "\n",
    "tk0 = tqdm(valid_loader)\n",
    "for i,(x_batch,)  in enumerate(tk0):\n",
    "    pred = model(x_batch.cuda(), attention_mask=(x_batch>0).cuda(), labels=None)\n",
    "    valid_preds[i*128:(i+1)*128]=pred[:,0].detach().cpu().squeeze().numpy()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate final score\n",
    "\n",
    "def calculate_overall_auc(df, model_name):\n",
    "    true_labels = df[TOXICITY_COLUMN]>0.5\n",
    "    predicted_labels = df[model_name]\n",
    "    return metrics.roc_auc_score(true_labels, predicted_labels)\n",
    "\n",
    "def power_mean(series, p):\n",
    "    total = sum(np.power(series, p))\n",
    "    return np.power(total / len(series), 1 / p)\n",
    "\n",
    "def get_final_metric(bias_df, overall_auc, POWER=-5, OVERALL_MODEL_WEIGHT=0.25):\n",
    "    bias_score = np.average([\n",
    "        power_mean(bias_df[SUBGROUP_AUC], POWER),\n",
    "        power_mean(bias_df[BPSN_AUC], POWER),\n",
    "        power_mean(bias_df[BNSP_AUC], POWER)\n",
    "    ])\n",
    "    return (OVERALL_MODEL_WEIGHT * overall_auc) + ((1 - OVERALL_MODEL_WEIGHT) * bias_score)\n",
    "\n",
    "\n",
    "\n",
    "SUBGROUP_AUC = 'subgroup_auc'\n",
    "BPSN_AUC = 'bpsn_auc'  # stands for background positive, subgroup negative\n",
    "BNSP_AUC = 'bnsp_auc'  # stands for background negative, subgroup positive\n",
    "\n",
    "def compute_auc(y_true, y_pred):\n",
    "    try:\n",
    "        return metrics.roc_auc_score(y_true, y_pred)\n",
    "    except ValueError:\n",
    "        return np.nan\n",
    "\n",
    "def compute_subgroup_auc(df, subgroup, label, model_name):\n",
    "    subgroup_examples = df[df[subgroup]>0.5]\n",
    "    return compute_auc((subgroup_examples[label]>0.5), subgroup_examples[model_name])\n",
    "\n",
    "def compute_bpsn_auc(df, subgroup, label, model_name):\n",
    "    \"\"\"Computes the AUC of the within-subgroup negative examples and the background positive examples.\"\"\"\n",
    "    subgroup_negative_examples = df[(df[subgroup]>0.5) & (df[label]<=0.5)]\n",
    "    non_subgroup_positive_examples = df[(df[subgroup]<=0.5) & (df[label]>0.5)]\n",
    "    examples = subgroup_negative_examples.append(non_subgroup_positive_examples)\n",
    "    return compute_auc(examples[label]>0.5, examples[model_name])\n",
    "\n",
    "def compute_bnsp_auc(df, subgroup, label, model_name):\n",
    "    \"\"\"Computes the AUC of the within-subgroup positive examples and the background negative examples.\"\"\"\n",
    "    subgroup_positive_examples = df[(df[subgroup]>0.5) & (df[label]>0.5)]\n",
    "    non_subgroup_negative_examples = df[(df[subgroup]<=0.5) & (df[label]<=0.5)]\n",
    "    examples = subgroup_positive_examples.append(non_subgroup_negative_examples)\n",
    "    return compute_auc(examples[label]>0.5, examples[model_name])\n",
    "\n",
    "def compute_bias_metrics_for_model(dataset,\n",
    "                                   subgroups,\n",
    "                                   model,\n",
    "                                   label_col,\n",
    "                                   include_asegs=False):\n",
    "    \"\"\"Computes per-subgroup metrics for all subgroups and one model.\"\"\"\n",
    "    records = []\n",
    "    for subgroup in subgroups:\n",
    "        record = {\n",
    "            'subgroup': subgroup,\n",
    "            'subgroup_size': len(dataset[dataset[subgroup]>0.5])\n",
    "        }\n",
    "        record[SUBGROUP_AUC] = compute_subgroup_auc(dataset, subgroup, label_col, model)\n",
    "        record[BPSN_AUC] = compute_bpsn_auc(dataset, subgroup, label_col, model)\n",
    "        record[BNSP_AUC] = compute_bnsp_auc(dataset, subgroup, label_col, model)\n",
    "        records.append(record)\n",
    "    return pd.DataFrame(records).sort_values('subgroup_auc', ascending=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>bnsp_auc</th>\n",
       "      <th>bpsn_auc</th>\n",
       "      <th>subgroup</th>\n",
       "      <th>subgroup_auc</th>\n",
       "      <th>subgroup_size</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.961452</td>\n",
       "      <td>0.906807</td>\n",
       "      <td>homosexual_gay_or_lesbian</td>\n",
       "      <td>0.863688</td>\n",
       "      <td>556</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0.967091</td>\n",
       "      <td>0.908245</td>\n",
       "      <td>white</td>\n",
       "      <td>0.880614</td>\n",
       "      <td>1269</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>0.960180</td>\n",
       "      <td>0.923157</td>\n",
       "      <td>muslim</td>\n",
       "      <td>0.882408</td>\n",
       "      <td>1049</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.973000</td>\n",
       "      <td>0.896655</td>\n",
       "      <td>black</td>\n",
       "      <td>0.893712</td>\n",
       "      <td>747</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.955299</td>\n",
       "      <td>0.944347</td>\n",
       "      <td>jewish</td>\n",
       "      <td>0.908925</td>\n",
       "      <td>379</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0.960492</td>\n",
       "      <td>0.942113</td>\n",
       "      <td>psychiatric_or_mental_illness</td>\n",
       "      <td>0.915371</td>\n",
       "      <td>235</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.942202</td>\n",
       "      <td>0.968407</td>\n",
       "      <td>christian</td>\n",
       "      <td>0.927764</td>\n",
       "      <td>2006</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.958788</td>\n",
       "      <td>0.960312</td>\n",
       "      <td>female</td>\n",
       "      <td>0.936172</td>\n",
       "      <td>2818</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.963767</td>\n",
       "      <td>0.954180</td>\n",
       "      <td>male</td>\n",
       "      <td>0.937014</td>\n",
       "      <td>2223</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   bnsp_auc  bpsn_auc                       subgroup  subgroup_auc  \\\n",
       "2  0.961452  0.906807      homosexual_gay_or_lesbian      0.863688   \n",
       "7  0.967091  0.908245                          white      0.880614   \n",
       "5  0.960180  0.923157                         muslim      0.882408   \n",
       "6  0.973000  0.896655                          black      0.893712   \n",
       "4  0.955299  0.944347                         jewish      0.908925   \n",
       "8  0.960492  0.942113  psychiatric_or_mental_illness      0.915371   \n",
       "3  0.942202  0.968407                      christian      0.927764   \n",
       "1  0.958788  0.960312                         female      0.936172   \n",
       "0  0.963767  0.954180                           male      0.937014   \n",
       "\n",
       "   subgroup_size  \n",
       "2            556  \n",
       "7           1269  \n",
       "5           1049  \n",
       "6            747  \n",
       "4            379  \n",
       "8            235  \n",
       "3           2006  \n",
       "1           2818  \n",
       "0           2223  "
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "0.9418549268955726"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "MODEL_NAME = 'model1'\n",
    "test_df[MODEL_NAME]=torch.sigmoid(torch.tensor(valid_preds)).numpy()\n",
    "TOXICITY_COLUMN = 'target'\n",
    "bias_metrics_df = compute_bias_metrics_for_model(test_df, identity_columns, MODEL_NAME, 'target')\n",
    "bias_metrics_df\n",
    "get_final_metric(bias_metrics_df, calculate_overall_auc(test_df, MODEL_NAME))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
